import asyncio
from typing import Iterator, AsyncIterator

from langchain_core.document_loaders import BaseLoader
from langchain_core.documents import Document


# ==============================
# 自定义文档加载器类
# ==============================
class CustomDocumentLoaddf(BaseLoader):
    """
    一个简单示例：从文件中逐行读取内容，生成 Document 对象
    """


    def __init__(self, file_path: str):
        """
        初始化加载器
        Args:
            file_path: 目标文件路径
        """
        self.file_path = file_path

    # ------------------------------
    # 同步惰性加载
    # ------------------------------
    def lazy_load(self) -> Iterator[Document]:
        """
        同步逐行读取文件，返回 Document 生成器
        - 这是“惰性加载”的实现方式，每次调用 yield 时只生成一条 Document
        - 优点：不占用大量内存，适合大文件
        """
        with open(self.file_path, encoding="utf-8") as f:
            line_number = 0
            for line in f:
                # yield 关键字用于定义生成器函数
                # 生成器函数式一种特殊类型的函数，它允许你逐步生成一个序列的值，而不是一次性返回整个序列。
                # 与普通的函数不同，生成器函数在每次调用时不会从头开始执行，而是从上次离开的地方继续执行。
                # 这使得生成器非常适合处理需要逐步生成或处理大数据集的情况
                yield Document(
                    page_content=line,
                    metadata={"line_number": line_number, "source": self.file_path},
                )
                line_number += 1

    # ------------------------------
    # 异步惰性加载 alazy_load 是可选的。如果省略了实现，将使用默认实现，该实现将委托给lazy_load
    # ------------------------------
    async def alazy_load(self) -> AsyncIterator[Document]:
        """
        异步逐行读取文件，返回异步生成器
        - 异步加载适用于 I/O 密集型操作（比如网络请求、文件读取）
        - 优点：在等待 I/O 时可以释放事件循环去处理其他任务
        """
        # 需要 aiofiles  使用pip install aiofiles 安装  https://github.com/Tinche/aiofiles
        import aiofiles
        async with aiofiles.open(self.file_path, encoding="utf-8") as f:
            line_number = 0
            async for line in f:
                # async for 逐行异步读取文件内容，并通过yield 关键字逐行生成Document文件
                # 异步生成器允许你再异步环境中逐步生成值，这对于处理I/O密集型任务（如文件读取或网络请求）非常有用
                # 因为它可以在登录I/O操作完成时释放事件循环，以便处理其他任务。
                yield Document(
                    page_content=line,
                    metadata={"line_number": line_number, "source": self.file_path},
                )
                line_number += 1


# ==============================
# 测试部分
# ==============================
# 1️⃣ 先创建一个测试文件
with open("/meow.txt","w",encoding="utf-8") as f:
    quality_context = "喵喵m \n 喵喵mm  \n 喵喵mmm"
    f.write(quality_context)

# 2️⃣ 初始化自定义加载器
loader = CustomDocumentLoaddf("/meow.txt")

# 3️⃣ 测试同步惰性加载
for doc in loader.lazy_load():
    print()
    print(type(doc))
    print(doc)

# 4️⃣ 测试异步惰性加载
async def alazy():
    async for doc in loader.alazy_load():
        print()
        print(type(doc))
        print(doc)

# asyncio.run 用于运行顶层异步函数
asyncio.run(alazy())
